import os
import sys
import django
import sqlite3
import pymysql
from datetime import datetime
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 设置Django环境
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'basal_adjust.settings')
django.setup()

print("开始迁移数据从SQLite到MySQL...")

# SQLite数据库路径
SQLITE_DB_PATH = 'db.sqlite3'

# MySQL连接配置 - 使用环境变量
MYSQL_CONFIG = {
    'host': os.environ.get('DB_HOST', 'localhost'),
    'user': os.environ.get('DB_USER', 'root'),
    'password': os.environ.get('DB_PASSWORD'),
    'database': os.environ.get('DB_NAME', 'basal_adjust'),
    'port': int(os.environ.get('DB_PORT', '3306')),
    'charset': 'utf8mb4'
}

def get_sqlite_tables():
    """获取SQLite数据库中所有表"""
    try:
        conn = sqlite3.connect(SQLITE_DB_PATH)
        cursor = conn.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
        tables = cursor.fetchall()
        conn.close()
        return [table[0] for table in tables if not table[0].startswith('sqlite_') and not table[0].startswith('django_')]
    except Exception as e:
        print(f"获取SQLite表失败: {e}")
        return []

def get_table_schema(table_name):
    """获取表结构"""
    try:
        conn = sqlite3.connect(SQLITE_DB_PATH)
        cursor = conn.cursor()
        cursor.execute(f"PRAGMA table_info({table_name});")
        columns = cursor.fetchall()
        conn.close()
        return columns
    except Exception as e:
        print(f"获取表 {table_name} 结构失败: {e}")
        return []

def get_table_data(table_name):
    """获取表数据"""
    try:
        conn = sqlite3.connect(SQLITE_DB_PATH)
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        cursor.execute(f"SELECT * FROM {table_name};")
        data = cursor.fetchall()
        conn.close()
        return data
    except Exception as e:
        print(f"获取表 {table_name} 数据失败: {e}")
        return []

def create_mysql_table(table_name, columns):
    """在MySQL中创建表"""
    try:
        # 检查数据库密码是否配置
        if not MYSQL_CONFIG['password']:
            print("错误：数据库密码未配置。请在.env文件中设置DB_PASSWORD")
            return False
            
        conn = pymysql.connect(**MYSQL_CONFIG)
        cursor = conn.cursor()
        
        # 创建表SQL
        column_defs = []
        primary_key = None
        
        for col in columns:
            col_id, name, col_type, notnull, default_value, pk = col
            
            # 根据SQLite类型映射到MySQL类型
            if col_type.upper() == 'INTEGER':
                mysql_type = 'INT'
            elif col_type.upper() == 'TEXT':
                mysql_type = 'TEXT'
            elif col_type.upper() == 'REAL':
                mysql_type = 'FLOAT'
            elif col_type.upper() == 'BLOB':
                mysql_type = 'BLOB'
            elif 'CHAR' in col_type.upper():
                mysql_type = col_type
            elif 'DATETIME' in col_type.upper():
                mysql_type = 'DATETIME'
            else:
                mysql_type = 'VARCHAR(255)'
            
            # 组装列定义
            col_def = f"`{name}` {mysql_type}"
            
            if notnull:
                col_def += " NOT NULL"
            
            if default_value is not None:
                if mysql_type in ('FLOAT', 'INT'):
                    col_def += f" DEFAULT {default_value}"
                else:
                    col_def += f" DEFAULT '{default_value}'"
                
            if pk:
                primary_key = name
                col_def += " AUTO_INCREMENT"
                
            column_defs.append(col_def)
        
        # 添加主键
        if primary_key:
            column_defs.append(f"PRIMARY KEY (`{primary_key}`)")
        
        # 创建表
        create_table_sql = f"CREATE TABLE IF NOT EXISTS `{table_name}` ({', '.join(column_defs)}) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;"
        
        print(f"创建表 {table_name} SQL: {create_table_sql}")
        cursor.execute(create_table_sql)
        conn.commit()
        cursor.close()
        conn.close()
        print(f"表 {table_name} 创建成功")
        return True
    except Exception as e:
        print(f"创建表 {table_name} 失败: {e}")
        return False

def insert_data_to_mysql(table_name, data, columns):
    """将数据插入到MySQL表"""
    if not data:
        print(f"表 {table_name} 没有数据，跳过")
        return
    
    try:
        # 检查数据库密码是否配置
        if not MYSQL_CONFIG['password']:
            print("错误：数据库密码未配置。请在.env文件中设置DB_PASSWORD")
            return
            
        conn = pymysql.connect(**MYSQL_CONFIG)
        cursor = conn.cursor()
        
        # 获取列名
        column_names = [col[1] for col in columns]
        
        # 准备SQL
        placeholders = ', '.join(['%s'] * len(column_names))
        insert_sql = f"INSERT INTO `{table_name}` (`{'`, `'.join(column_names)}`) VALUES ({placeholders})"
        
        # 批量插入数据
        batch_size = 100
        for i in range(0, len(data), batch_size):
            batch_data = []
            for row in data[i:i+batch_size]:
                row_data = []
                for name in column_names:
                    val = row[name]
                    if isinstance(val, bytes):
                        try:
                            val = val.decode('utf-8')
                        except UnicodeDecodeError:
                            val = str(val)
                    row_data.append(val)
                batch_data.append(row_data)
            
            cursor.executemany(insert_sql, batch_data)
            conn.commit()
            print(f"表 {table_name} 插入数据: {len(batch_data)} 行")
        
        cursor.close()
        conn.close()
        print(f"表 {table_name} 数据迁移完成，共 {len(data)} 行")
    except Exception as e:
        print(f"表 {table_name} 数据插入失败: {e}")

def check_mysql_table_exists(table_name):
    """检查MySQL表是否存在"""
    try:
        # 检查数据库密码是否配置
        if not MYSQL_CONFIG['password']:
            print("错误：数据库密码未配置。请在.env文件中设置DB_PASSWORD")
            return False
            
        conn = pymysql.connect(**MYSQL_CONFIG)
        cursor = conn.cursor()
        cursor.execute(f"SHOW TABLES LIKE '{table_name}';")
        result = cursor.fetchone()
        cursor.close()
        conn.close()
        return result is not None
    except Exception as e:
        print(f"检查表 {table_name} 是否存在失败: {e}")
        return False

def migrate_table(table_name):
    """迁移一个表"""
    print(f"\n开始迁移表: {table_name}")
    
    # 检查MySQL表是否已存在
    if check_mysql_table_exists(table_name):
        print(f"表 {table_name} 在MySQL中已存在，跳过创建")
        return
    
    # 获取表结构
    columns = get_table_schema(table_name)
    if not columns:
        print(f"无法获取表 {table_name} 结构，跳过迁移")
        return
    
    # 创建MySQL表
    if create_mysql_table(table_name, columns):
        # 获取数据并插入
        data = get_table_data(table_name)
        if data:
            insert_data_to_mysql(table_name, data, columns)
        else:
            print(f"表 {table_name} 没有数据")

def main():
    """主函数"""
    # 检查环境变量配置
    if not MYSQL_CONFIG['password']:
        print("错误：数据库密码未配置")
        print("请执行以下步骤：")
        print("1. 复制 env.example 为 .env")
        print("2. 在 .env 文件中设置 DB_PASSWORD=你的数据库密码")
        return
    
    print(f"连接到数据库: {MYSQL_CONFIG['host']}:{MYSQL_CONFIG['port']}/{MYSQL_CONFIG['database']}")
    
    tables = get_sqlite_tables()
    if not tables:
        print("没有找到需要迁移的表")
        return
    
    print(f"找到 {len(tables)} 个表需要迁移: {', '.join(tables)}")
    
    for table_name in tables:
        migrate_table(table_name)
    
    print("\n数据迁移完成！")

if __name__ == "__main__":
    main() 